from PIL import Image
from torchvision import transforms
from torch import fft
import matplotlib.pyplot as plt
import torch

transform_ = transforms.Compose([
    transforms.ToTensor()
])

img = Image.open("./bottle.png").convert('L')
img_tensor = transform_(img)
img_tensor = torch.cat((img_tensor, img_tensor), dim=0)
print(img_tensor.shape)
dft = fft.fftn(img_tensor, dim=(-2, -1))

magnitude_spectrum = torch.log(torch.abs(dft) + 1)

plt.subplot(121), plt.imshow(img_tensor.view(-1, 900).numpy(), cmap='gray')
plt.title('Input Image'), plt.xticks([]), plt.yticks([])
plt.subplot(122), plt.imshow(magnitude_spectrum.view(-1, 900).numpy(), cmap='gray')
plt.title('Magnitude Spectrum'), plt.xticks([]), plt.yticks([])
plt.show()
